from flask import Flask, render_template, request, jsonify
import os
import sys
import tempfile
import shutil
from werkzeug.utils import secure_filename

# 添加当前目录到Python路径
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

from ask import LocalQASystem

app = Flask(__name__)
app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024  # 最大上传16MB
app.config['UPLOAD_FOLDER'] = tempfile.gettempdir()

# 初始化问答系统
qa_system = LocalQASystem()
print("正在初始化问答系统...")
if not qa_system.init_system():
    print("系统初始化失败！")
    sys.exit(1)

@app.route('/')
def index():
    return render_template('index.html')

@app.route('/ask', methods=['POST'])
def ask():
    try:
        data = request.get_json()
        question = data.get('question', '').strip()
        
        if not question:
            return jsonify({'error': '问题不能为空'}), 400
            
        answer = qa_system.ask_question(question)
        return jsonify({'answer': answer})
        
    except Exception as e:
        return jsonify({'error': str(e)}), 500

@app.route('/upload', methods=['POST'])
def upload_file():
    try:
        if 'file' not in request.files:
            return jsonify({'error': '没有文件'}), 400
            
        file = request.files['file']
        if file.filename == '':
            return jsonify({'error': '没有选择文件'}), 400
            
        if file and file.filename.endswith('.txt'):
            filename = secure_filename(file.filename)
            
            # 获取项目根目录
            cur_dir = os.path.dirname(os.path.abspath(__file__))
            data_dir = os.path.join(cur_dir, 'data')
            os.makedirs(data_dir, exist_ok=True)
            
            # 保存文件到data目录
            filepath = os.path.join(data_dir, filename)
            file.save(filepath)
            
            # 验证文件是否保存成功
            if not os.path.exists(filepath):
                return jsonify({'error': '文件保存失败'}), 500
            
            # 重新初始化系统以包含新文档
            print(f"正在重新初始化系统以包含新文档: {filename}")
            if qa_system.init_system():
                return jsonify({'message': f'文件 {filename} 上传成功并已添加到知识库'})
            else:
                # 如果初始化失败，删除刚上传的文件
                try:
                    os.remove(filepath)
                except:
                    pass
                return jsonify({'error': '文件上传成功，但系统重新初始化失败'}), 500
        else:
            return jsonify({'error': '只支持txt文件'}), 400
            
    except Exception as e:
        return jsonify({'error': f'上传失败: {str(e)}'}), 500

@app.route('/database/info')
def database_info():
    try:
        # 获取分页参数
        page = int(request.args.get('page', 1))
        per_page = int(request.args.get('per_page', 5))
        
        if qa_system.vectordb is None:
            return jsonify({'error': '数据库未初始化'}), 400
            
        count = qa_system.vectordb._collection.count()
        
        # 计算总页数
        total_pages = (count + per_page - 1) // per_page if count > 0 else 1
        page = max(1, min(page, total_pages))
        
        # 计算偏移量
        offset = (page - 1) * per_page
        
        # 获取指定页面的文档
        docs = qa_system.vectordb._collection.get(limit=per_page, offset=offset)
        
        doc_previews = []
        for i in range(len(docs['ids'])):
            preview = {
                'id': docs['ids'][i],
                'content': docs['documents'][i][:200] + '...' if len(docs['documents'][i]) > 200 else docs['documents'][i],
                'metadata': docs['metadatas'][i] if i < len(docs['metadatas']) else {}
            }
            doc_previews.append(preview)
            
        return jsonify({
            'document_count': count,
            'current_page': page,
            'total_pages': total_pages,
            'per_page': per_page,
            'document_previews': doc_previews
        })
        
    except Exception as e:
        return jsonify({'error': str(e)}), 500

@app.route('/system/status')
def system_status():
    try:
        status = qa_system.get_system_status()
        return jsonify(status)
    except Exception as e:
        return jsonify({'error': str(e)}), 500

# 添加一个手动重新初始化系统的接口
@app.route('/reinit', methods=['POST'])
def reinit_system():
    try:
        if qa_system.init_system():
            return jsonify({'message': '系统重新初始化成功'})
        else:
            return jsonify({'error': '系统重新初始化失败'}), 500
    except Exception as e:
        return jsonify({'error': str(e)}), 500

if __name__ == '__main__':
    app.run(debug=True, host='0.0.0.0', port=5000)